# -*- coding: utf-8 -*-
#!/usr/bin/env python

''' Autor: Neuza Figueira - n.º 6036 '''

from math import sqrt
import sys
import time
import random
import pylab
from random import choice

#Class to create a KDTree and manage it's operations.
class KDTree:
    #Constructor.
    #@param points -> list of values to create the nodes.
    #@param N -> number of nodes, to deal with the memory allocation.
    def __init__(self, points, N ):
            self.N = N
            self.clear()
            self.root = self.create(points)
            pass

	
    # To create the tree.
    # @param points -> list of values to create the nodes.
    # @param depth -> height of the tree.
    def create(self, points, depth = 0):
                    
            if not points:
                    return -1

            x = len(points[0])
            a = depth % x

            points.sort(key = lambda point: point[a])
            m = len(points) / 2

            i = self.malloc(points[m])

            self.left[i] = self.create(points[:m], depth + 1)
            self.right[i] = self.create(points[m + 1:], depth + 1)

            return i

        
    # To add new nodes to the tree.
    #@ param z -> coordinates of the new node.
    def addNode(self, z):
            if not self.pointer.stack_full():
                    points = []
                    for k in xrange(0, self.N):
                            if self.z[k] != []:
                                    points.append(self.z[k])
                                    pass
                            
                            pass

                    points.append(z)
                    self.clear()
                    
                    if points:
                            self.root = self.create(points)
                    pass


    # To clear the tree.
    def clear(self):
            self.pointer = Stack(self.N)
            for x in xrange(self.N):
                    self.pointer.push(x)
            self.z =  [[] for x in xrange(self.N)]
            self.parent = [-1 for x in xrange(self.N)]
            self.left = [-1 for x in xrange(self.N)]
            self.right = [-1 for x in xrange (self.N)]
            self.root = -1
            pass

    # To remove a node from the tree.
    # @param z -> coordinates of the node to delete.
    def delete(self, z):
            p = self.search(z)
            if p != -1:
                    points = []
                    for k in xrange(0, self.N):
                            if self.z[k] != [] and self.z[k] != z:
                                    points.append(self.z[k])
                                    pass
                            
                            pass

                    self.clear()
                    if points:
                            self.root = self.create(points)
                    pass


    #To calculate the distance between nodes.
    # @param z -> coordinates of the node.
    # @param x
    def getDistance(self, z, x):
            distance = 0
            s = len (z)
            for k in xrange(0 ,s):
                    temp = self.z[x][k]-z[k] 
                    distance += temp * temp
                    pass
            distance = sqrt(distance)
            return distance

    # To apply the Nearest Neighbour Algorithm.
    # @param z -> coordinates of the node to search the nearest neighbour.
    def nearestNeighborSearch(self, z):
        # To add all the nodes in the chosen path to the stack.
        # @param kd -> the KDTree object.
        # @param z -> coordinates of the node to search the nearest neighbour.
        # @param node -> the node to search the nearest neighbour.
        # @param a -> stands for the axis
        # @param stack -> the stack object
            def leafSearch (kd, z, node, a, stack):
                    i = node                    
                    s = len (z)
                    a = a % s
                    
                    while i != -1 and kd.z[i] != z:
                            if kd.nodeComparison(a, z, i):
                                    i = kd.right[i]
                            else:
                                    i = kd.left[i]

                            if i != -1:
                                    stack.push((i, a))

                            a = (a + 1) % s
                            
                            pass
                    pass

            d = sys.maxint
            stack = Stack(self.N)

            node = self.root
            a = 0

            leafSearch(self, z, node, a, stack)

            last_node = -1

            while not stack.stack_empty():
                    nodeA = stack.pop()
                    
                    distanceTemp = self.getDistance(z, nodeA[0])
                    distance = sys.maxint

                    if distanceTemp < distance:
                            distance = distanceTemp
                            node = nodeA[0]
                            pass

                    axisDistance = abs(z[nodeA[1]] - self.z[nodeA[0]][nodeA[1]])

                    if axisDistance < distance:
                            if last_node == self.right[nodeA[0]]:
                                    leafSearch(self, z, self.left[nodeA[0]], nodeA[1] + 1, stack)
                            else:
                                    leafSearch(self, z, self.right[nodeA[0]], nodeA[1] + 1, stack)
                                    pass
                            pass


                    last_node = nodeA[0]
                    pass

            return node

   # To alloc memory.
   # @param k
    def malloc(self, k):
            x = self.pointer.pop()
            self.z[x] = k
            self.parent[x] = -1
            self.left[x] = -1
            self.right[x] = -1
            return x


    # To search a given point.
    # @param z -> coordinates of the node to search.
    def search (self, z):
            index = self.root
            a = 0

            s = len (z)

            while index != -1 and self.z[index] != z:
                    if self.nodeComparison(a, z, index):
                            index = self.right[index]
                    else:
                            index = self.left[index]

                    a = (a + 1) % s


            return index


    # To compare nodes.
    # @param a -> stands for the axis.
    # @param z -> the coordinates of the node.
    # @param node -> the node.
    def nodeComparison(self, a, z, node):
            s = len (z) 

            for i in xrange(a, s + a):
                    index = i % s
                    if z[index] > self.z[node][index]:
                            return True
                    elif z[index] < self.z[node][index]:
                            return False
                            pass
                    pass

            return False

    # To print information of a given index.
    # @param index
    def printNode (self, index):
            s = '['

            if self.left[index] == -1:
                    s += 'None'
            else:
                    s += self.printNode(self.left[index])

            s +=  ', '

            if index == self.root:
                    s+= 'root: '
            
            s += str(self.z[index]) + ', '

            if self.right[index] == -1:
                    s += 'None'
            else:
                    s+=self.printNode(self.right[index])

            s+=']'

            return s

	

    # To print the whole tree.
    def __str__ (self):
            s = '['

            if self.left[self.root] == -1:
                    s += 'None'
            else:
                    s += self.printNode(self.left[self.root])

            s +=  ', '

            if self.root == self.root:
                    s+= 'root: '
            
            s += str(self.z[self.root]) + ', '

            if self.right[self.root] == -1:
                    s += 'None'
            else:
                    s+=self.printNode(self.right[self.root])

            s+=']'

            return s
        
# Clas to create the stack.
class Stack:
        def __init__(self, N):
                self.top = -1
                self.S = [0 for k in xrange(N)]
                pass
        
        # To check if the stack is empty.
        def stack_empty(self):
                if self.top < 0:
                        return True
                else:
                        return False
        
        # To check if the stack is full.
        def stack_full(self):
                if self.top >= len (self.S) -1:
                        return True
                else:
                        return False
                pass
                

        # To put values in the stack.
        def push(self, x):
                self.top += 1
                self.S[self.top] = x
                pass
        
        # To take values from the stack.
        def pop(self):
                if self.stack_empty():
                        return "underflow"
                else:
                        self.top -= 1
                        return self.S[self.top + 1]
                pass


# To make tests.
if __name__ == "__main__":
    
    listaKD = [(1,0,2,3), (5,6,7,4), (7,8,4,2), (3,9,5,3), (3,5,7,9), (0,2,3,7), (1,3,5,6), (9,7,4,2), (0,5,2,3)]
    kd = KDTree(listaKD, 10)
    print "_____________________________________KDTree_____________________________________"
    print kd
    
    print "___________________________________Add a node___________________________________"    
    kd.addNode((4,5,6,1))
    print kd

    print "__________________________________Delete a node_________________________________"
    kd.delete((7,8,4,2))
    kd.delete((3,5,7,9))
    print kd

    print "________________________________Nearest Neighbour_______________________________"
    N = kd.nearestNeighborSearch((1,0,2,3))
    print kd.printNode(N)

    # Time tests.
    ''' Test One'''
    print "____________________________________Tests_______________________________________"
    start = 0
    maxN = 50
    maxValue = 10000
    lista = [None for k in range(maxN)]

    for j in range(maxN):
        lista[j] = [random.randint(start, maxValue), random.randint(start, maxValue)]
        pass

    t1 = time.clock()
    tree = KDTree(lista,  maxN)
    t2 = time.clock()

    print "Test One - Insertion Time:"
    print t2-t1 / 60

    t3 = time.clock()
    N = kd.nearestNeighborSearch(choice(lista))
    t4= time.clock()
    print "Test One - Nearest Neighbour:"
    print kd.printNode(N)

    print "Test One - Search Time:"
    print t4-t3 / 60
    print "_______________________________________________________________________________"
    
    ''' Test Two'''
    start = 0
    maxN = 500
    maxValue = 10000
    lista = [None for k in range(maxN)]

    for j in range(maxN):
        lista[j] = [random.randint(start, maxValue), random.randint(start, maxValue), random.randint(start, maxValue)]
        pass

    t1 = time.clock()
    tree = KDTree(lista,  maxN)
    t2 = time.clock()

    print "Test Two - Insertion Time:"
    print t2-t1 / 60
    
    t3 = time.clock()
    N = kd.nearestNeighborSearch(choice(lista))
    t4 = time.clock()
    print "Test Two - Nearest Neighbour:"
    print kd.printNode(N)
    
    print "Test Two - Search Time:"
    print t4-t3 / 60
    print "_______________________________________________________________________________"
    
    ''' Test Three'''
    start = 0
    maxN = 5000
    maxValue = 10000
    lista = [None for k in range(maxN)]

    for j in range(maxN):
        lista[j] = [random.randint(start, maxValue),random.randint(start, maxValue), random.randint(start, maxValue), random.randint(start, maxValue)]
        pass

    t1 = time.clock()
    tree = KDTree(lista,  maxN)
    t2 = time.clock()

    print "Test Three - Insertion Time:"
    print t2-t1 / 60

    t3 = time.clock()
    N = kd.nearestNeighborSearch(choice(lista))
    t4 = time.clock()
    print "Test Three - Nearest Neighbour:"
    print kd.printNode(N)

    print "Test Three - Search Time:"
    print t4-t3 / 60
    print "_______________________________________________________________________________"



    

    
    

    




